# main.py
# 该脚本用于编排整个考试自动答题和书写的流水线。
# 它会按顺序调用四个核心功能的脚本，并把所有控制台输出实时记录到一个总日志文件中。

import os
import sys
import traceback
import datetime

# --- 核心改动：定义一个日志记录类 ---
class Logger:
    """
    一个特殊的日志记录类。
    它可以将所有写入操作（如 print）同时输出到原始控制台和一个日志文件中。
    """
    def __init__(self, log_filepath, original_stdout):
        self.original_stdout = original_stdout  # 保存原始的控制台输出对象
        self.log_file = open(log_filepath, "a", encoding="utf-8")  # 以追加模式打开日志文件

    def write(self, message):
        """
        重写 write 方法，实现双重输出。
        """
        self.original_stdout.write(message)  # 先在控制台打印
        self.log_file.write(message)         # 然后写入到文件
        self.flush()  # 每次写入后立即刷新，确保内容被实时保存

    def flush(self):
        """
        重写 flush 方法，确保两个输出流都被刷新。
        """
        self.original_stdout.flush()
        self.log_file.flush()
    
    def close(self):
        """
        关闭文件句柄。
        """
        self.log_file.close()

def run_pipeline():
    """
    通过按正确顺序导入并执行每个脚本的 main 函数，来运行完整的流水线。
    （这个函数内部的逻辑和之前一样，无需改动）
    """
    print("=====================================================")
    print("===         AI 驱动的考试答题流水线         ===")
    print("=====================================================")

    # 将当前脚本的目录添加到 Python 路径。
    script_dir = os.path.dirname(os.path.abspath(__file__))
    if script_dir not in sys.path:
        sys.path.append(script_dir)

    # --- 第 1 步：捕获、旋转和 OCR 页面 ---
    try:
        print("\n[第 1/4 步] 初始化页面捕获与 OCR...")
        import capture
        capture.main()
        print("\n[成功] 第 1 步：页面捕获与 OCR 已成功完成。")
    except Exception as e:
        print(f"\n[致命错误] 在第 1 步 (capture.py) 中发生错误：")
        traceback.print_exc()
        print("\n流水线已停止。")
        return

    # --- 第 2 步：从 OCR 文本生成答案 ---
    try:
        print("\n[第 2/4 步] 正在从 OCR 文本生成答案...")
        import answer
        answer.main()
        print("\n[成功] 第 2 步：答案生成已成功完成。")
    except Exception as e:
        print(f"\n[致命错误] 在第 2 步 (answer.py) 中发生错误：")
        traceback.print_exc()
        print("\n流水线已停止。")
        return

    # --- 第 3 步：将答案对齐到红色方框位置 ---
    try:
        print("\n[第 3/4 步] 正在将答案映射到每页的红色方框位置...")
        import align_positions
        align_positions.main()
        print("\n[成功] 第 3 步：答案到位置的映射已成功完成。")
    except Exception as e:
        print(f"\n[致命错误] 在第 3 步 (align_positions.py) 中发生错误：")
        traceback.print_exc()
        print("\n流水线已停止。")
        return

    # --- 第 4 步：规划并执行物理书写 ---
    try:
        print("\n[第 4/4 步] 正在规划书写任务并由机械臂执行...")
        import document_processor
        document_processor.main()
        print("\n[成功] 第 4 步：文档处理与书写已成功完成。")
    except Exception as e:
        print(f"\n[致命错误] 在第 4 步 (document_processor.py) 中发生错误：")
        traceback.print_exc()
        print("\n流水线已停止。")
        return

    print("\n=====================================================")
    print("===         🎉 流水线全部执行成功！      ===")
    print("=====================================================")


if __name__ == "__main__":
    # --- 核心改动：设置日志文件路径并重定向标准输出 ---
    LOG_FILE_PATH = r"E:\沙粒云\自媒体\2025视频制作\20250610ai高考\流程\anwser\data\output.txt"
    
    # 确保日志文件所在的目录存在
    os.makedirs(os.path.dirname(LOG_FILE_PATH), exist_ok=True)
    
    # 每次运行时，先清空（或创建）日志文件，并写入一个启动时间戳
    with open(LOG_FILE_PATH, "w", encoding="utf-8") as f:
        f.write(f"流水线启动于 {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
        f.write("="*50 + "\n\n")

    # 保存原始的 sys.stdout
    original_stdout = sys.stdout
    
    # 创建我们的 Logger 实例，并用它替换掉 sys.stdout
    logger_instance = Logger(LOG_FILE_PATH, original_stdout)
    sys.stdout = logger_instance

    try:
        # 现在，所有对 print 的调用都会通过我们的 Logger
        run_pipeline()
    finally:
        # 无论流水线是否成功，最后都必须恢复原始的 stdout 并关闭文件
        sys.stdout = original_stdout
        logger_instance.close()
        print("总日志文件已关闭。")